# -*- coding: utf-8 -*-
# $ ./build/triton ./src/examples/pin/inject_model_with_snapshot.py ./src/samples/crackmes/crackme_xor a
from triton import ARCH
from pintool import *
import sys

password = dict()
cur_char_ptr = None
Triton = getTritonContext()


def csym(instruction):
    global cur_char_ptr

    # print(instruction)

    # 目标函数的入口地址， 第一次执行就拍个快照
    if instruction.getAddress() == 0x400556 and isSnapshotEnabled() == False:
        takeSnapshot()
        return
    if instruction.getAddress() == 0x400574:
        rax = getCurrentRegisterValue(Triton.registers.rax)
        cur_char_ptr = rax  # 每次取字符的地址

        # 如果这个位置已经求出解了，就设置解
        if rax in password:
            setCurrentMemoryValue(rax, password[rax])
        return

    # check 函数的结尾, 判断返回值是否满足要求
    if instruction.getAddress() == 0x4005b2:
        rax = getCurrentRegisterValue(Triton.registers.rax)
        # 如果 rax 不是 0 ， 说明还需要继续求解
        if rax != 0:
            # 恢复快照，继续运行
            restoreSnapshot()
        else:
            disableSnapshot()
            # 把解由地址从低往高开始打印
            addrs = password.keys()
            addrs.sort()
            answer = ""
            for addr in addrs:
                c = chr(password[addr])
                answer += c
                print("0x{:08x}: {}".format(addr, c))
            print("answer: {}".format(answer))
        return
    return

def cafter(instruction):
    global password
    # print(instruction)

    # 0000400574 movzx   eax, byte ptr [rax]
    # 执行完 0x400574 后， rax 里面存放的是刚刚从输入字符串中取出的字符
    # 将取出的字符设置为符号量， 后面用来求解
    if instruction.getAddress() == 0x400574:
        var = Triton.convertRegisterToSymbolicVariable(Triton.registers.rax)
        return

    # 400597 cmp     ecx, eax
    # 开始求解约束
    if instruction.getAddress() == 0x400597:
        astCtxt = Triton.getAstContext()
        zf = Triton.getRegisterAst(Triton.registers.zf)
        # 如果正确的话，需要 zf == 1， 增加约束
        cstr = astCtxt.land([
            Triton.getPathConstraintsAst(),
            zf == 1
        ])
        models = Triton.getModel(cstr)
        for k, v in list(models.items()):
            # 把计算的结果保存
            password.update({cur_char_ptr: v.getValue()})
        return
    return


def fini():
    print('[+] Analysis done!')
    return




if __name__ == '__main__':
    setupImageWhitelist(['crackme_xor'])

    # 设置从 check 函数开始分析
    # startAnalysisFromSymbol('check')

    startAnalysisFromAddress(0x0400556)

    
    insertCall(cafter,  INSERT_POINT.AFTER)
    insertCall(csym,    INSERT_POINT.BEFORE_SYMPROC)

    insertCall(fini,    INSERT_POINT.FINI)
    runProgram()
